#!/usr/bin/env python3
"""
使用分割后的数据进行DexiNed训练的脚本
"""

import os
import sys
import argparse
import subprocess

def train_with_split_data():
    """使用分割后的数据进行训练"""
    
    # 数据路径配置
    split_data_dir = "/home/dd/working/data/dexined/cropped_hy_images_split"
    train_list = "train_pair.lst"
    val_list = "val_pair.lst"
    
    # 检查数据是否存在
    if not os.path.exists(split_data_dir):
        print(f"错误: 分割后的数据目录不存在: {split_data_dir}")
        print("请先运行 split_dataset.py 进行数据分割")
        return False
    
    if not os.path.exists(os.path.join(split_data_dir, train_list)):
        print(f"错误: 训练列表文件不存在: {os.path.join(split_data_dir, train_list)}")
        return False
    
    if not os.path.exists(os.path.join(split_data_dir, val_list)):
        print(f"错误: 验证列表文件不存在: {os.path.join(split_data_dir, val_list)}")
        return False
    
    # 训练参数
    training_args = [
        "python", "main.py",
        "--train_data", "BSDS",
        "--test_data", "BSDS", 
        "--input_dir", split_data_dir,
        "--input_val_dir", split_data_dir,
        "--train_list", train_list,
        "--test_list", val_list,
        "--output_dir", "checkpoints",
        "--epochs", "20",
        "--batch_size", "4",
        "--lr", "1e-4",
        "--wd", "1e-8",
        "--img_width", "352",
        "--img_height", "352",
        "--test_img_width", "352", 
        "--test_img_height", "352",
        "--workers", "2",
        "--log_interval_vis", "20",
        "--crop_img", "True",
        "--tensorboard", "True"
    ]
    
    print("开始训练...")
    print("训练命令:")
    print(" ".join(training_args))
    print()
    
    # 执行训练
    try:
        result = subprocess.run(training_args, check=True)
        print("训练完成!")
        return True
    except subprocess.CalledProcessError as e:
        print(f"训练过程中出现错误: {e}")
        return False
    except KeyboardInterrupt:
        print("训练被用户中断")
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='训练DexiNed模型使用分割后的数据')
    parser.add_argument('--dry_run', action='store_true', 
                       help='只显示命令不执行训练')
    
    args = parser.parse_args()
    
    if args.dry_run:
        print("模拟运行模式，只显示训练命令:")
        print()
        train_with_split_data()
    else:
        success = train_with_split_data()
        if success:
            print("\\n训练成功完成!")
        else:
            print("\\n训练失败!")
            sys.exit(1)